/*
  Copyright 2011 Keith E. Laidig

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
 */
package org.aimj.tools;
import java.io.*;
import java.text.*;
import java.util.*;
import org.aimj.atomicfunctions.*;
/*
 * WFN.class
 * KEL
 * a class that represents the wavefunction
 */
public class WFN {
    public String titleWFN;  // wavefunction title
    public Integer mode, numOrbitals, numPrimitives, numCentres;  // mode, number of orbitals, number of primitives, number of function centres (atoms)
    public List<Centre> XYZ = new ArrayList<Centre>();  // atomic coordinates
    public List<Double> charge = new ArrayList<Double>();  // nuclear charge
    public List<AF> atomicFunctions = new ArrayList<AF>();  // atomic functions
    public List<Double> orbitalOccupation = new ArrayList<Double>();  // orbital occupation
    public List<Double> orbitalEnergy = new ArrayList<Double>();  // orbital energy
    public Double totE, gamma;  // total energy, virial ratio

    /*
     * class constructor
     */
    public WFN(BufferedReader br) throws ParseException, IOException {
	 		readWFN(br);  // call method to load up class member variables
    }
    /*
     * read in data from WFN file
     * use FortranFormat class to read in formatted data from BufferedReader
     */
    public void readWFN(BufferedReader br) throws ParseException, IOException {
        // temp ArrayLists
        List<Integer> functionCentre = new ArrayList<Integer>();  // temp function centres
        List<Integer> functionTypes = new ArrayList<Integer>();  // temp function types
        //
        //READ(IWFN,101) WFNTTL
        titleWFN = br.readLine();
        String line = br.readLine();
        // read in wavefunction parameters
        //READ(IWFN,102) MODE,NMO,NPRIMS,NCENT
        List<Object> objects = FortranFormat.read(line,"(4X,A4,12X,3(I3,17X))");
        numOrbitals = (Integer)objects.get(1);
        numPrimitives = (Integer)objects.get(2);
        numCentres = (Integer)objects.get(3);
        // read in atomic parameters
        //DO 100 I = 1,NCENT
        //READ (IWFN,103) ATNAM(I),J,CO(IXC+J),CO(IYC+J),CO(IZC+J),CO(ICHARG+J)
        for (int i = 0 ; i < numCentres ; ++i) {
            line = br.readLine();
            objects = FortranFormat.read(line,"(A8,11X,I3,2X,3F12.8,10X,F5.1)");
            XYZ.add(new Centre((Double)objects.get(2),(Double)objects.get(3), (Double)objects.get(4)));
            charge.add((Double)objects.get(5));
        }
        // read in function centres
        //READ (IWFN,104) (IC(ICENT+I),I=1,NPRIMS)
        line = br.readLine();
        while (line.startsWith("CENTRE ASSIGNMENTS")) {
            objects = FortranFormat.read(line,"(20X,20I3)");
            for (int i = 0; i < 20; ++i) {  // 20 objects per line
                 if (objects.get(i) != null)
                    functionCentre.add((Integer)objects.get(i));
            }
            line = br.readLine();  // read next line
        }
        //System.out.println(numPrimitives + " primitives and " + functionCentres.size() + " function centres");
        // read in function types
        //READ (IWFN,104) (IC(ITYPE+I),I=1,NPRIMS)
        while (line.startsWith("TYPE ASSIGNMENTS")) {
            objects = FortranFormat.read(line,"(20X,20I3)");
            for (int i = 0; i < 20; ++i) {  // 20 objects per line
                if (objects.get(i) != null)
                    functionTypes.add((Integer)objects.get(i));
            }
            line = br.readLine();  // read next line
        }
        //System.out.println(numPrimitives + " primitives and " + functionTypes.size() + " function types");
        // read in function coefficients
        //READ (IWFN,105) (CO(IE+I),I=1,NPRIMS)
        while (line.startsWith("EXPONENTS")) {
            objects = FortranFormat.read(line,"(10X,5E14.7)");
            for (int i = 0; i < 5; ++i) {  // 5 objects per line
                if (objects.get(i) != null) {
                    // here we create the atomic functions given function type....
                    if (functionTypes.get(i) == 1) {
                        atomicFunctions.add(new sAF(XYZ.get(functionCentre.get(i) - 1), (Double)objects.get(i)));  // create an s-type function
                    } else if (functionTypes.get(i) == 2) {
                        atomicFunctions.add(new pAF(XYZ.get(functionCentre.get(i) - 1), (Double)objects.get(i), 1));  // create a px-type function
                    } else if (functionTypes.get(i) == 3) {
                        atomicFunctions.add(new pAF(XYZ.get(functionCentre.get(i) - 1), (Double)objects.get(i), 2));  // create a py-type function
                    } else if (functionTypes.get(i) == 4) {
                        atomicFunctions.add(new pAF(XYZ.get(functionCentre.get(i) - 1), (Double)objects.get(i), 3));  // create a pz-type function
                    } else if (functionTypes.get(i) == 5) {
                        atomicFunctions.add(new dAF(XYZ.get(functionCentre.get(i) - 1), (Double)objects.get(i), 1, 1));  // create a dxx-type function
                    } else if (functionTypes.get(i) == 6) {
                        atomicFunctions.add(new dAF(XYZ.get(functionCentre.get(i) - 1), (Double)objects.get(i), 1, 2));  // create a dxy-type function
                    } else if (functionTypes.get(i) == 7) {
                        atomicFunctions.add(new dAF(XYZ.get(functionCentre.get(i) - 1), (Double)objects.get(i), 1, 3));  // create a dxz-type function
                    } else if (functionTypes.get(i) == 8) {
                        atomicFunctions.add(new dAF(XYZ.get(functionCentre.get(i) - 1), (Double)objects.get(i), 2, 2));  // create a dyy-type function
                    } else if (functionTypes.get(i) == 9) {
                        atomicFunctions.add(new dAF(XYZ.get(functionCentre.get(i) - 1), (Double)objects.get(i), 2, 3));  // create a dyz-type function
                    } else if (functionTypes.get(i) == 10) {
                        atomicFunctions.add(new dAF(XYZ.get(functionCentre.get(i) - 1), (Double)objects.get(i), 3, 3));  // create a dzz-type function
                    } else if (functionTypes.get(i) == 11) {
                        atomicFunctions.add(new fAF(XYZ.get(functionCentre.get(i) - 1), (Double)objects.get(i), 1, 1, 1));  // create an fxxx-type function
                    } else if (functionTypes.get(i) == 12) {
                        atomicFunctions.add(new fAF(XYZ.get(functionCentre.get(i) - 1), (Double)objects.get(i), 1, 1, 2));  // create an fxxy-type function
                    } else if (functionTypes.get(i) == 13) {
                        atomicFunctions.add(new fAF(XYZ.get(functionCentre.get(i) - 1), (Double)objects.get(i), 1, 1, 3));  // create an fxxz-type function
                    } else if (functionTypes.get(i) == 14) {
                        atomicFunctions.add(new fAF(XYZ.get(functionCentre.get(i) - 1), (Double)objects.get(i), 1, 2, 2));  // create an fxyy-type function
                    } else if (functionTypes.get(i) == 15) {
                        atomicFunctions.add(new fAF(XYZ.get(functionCentre.get(i) - 1), (Double)objects.get(i), 1, 2, 3));  // create an fxyz-type function
                    } else if (functionTypes.get(i) == 16) {
                        atomicFunctions.add(new fAF(XYZ.get(functionCentre.get(i) - 1), (Double)objects.get(i), 1, 3, 3));  // create an fxzz-type function
                    } else if (functionTypes.get(i) == 17) {
                        atomicFunctions.add(new fAF(XYZ.get(functionCentre.get(i) - 1), (Double)objects.get(i), 2, 2, 2));  // create an fyyy-type function
                    } else if (functionTypes.get(i) == 18) {
                        atomicFunctions.add(new fAF(XYZ.get(functionCentre.get(i) - 1), (Double)objects.get(i), 2, 2, 3));  // create an fyyz-type function
                    } else if (functionTypes.get(i) == 19) {
                        atomicFunctions.add(new fAF(XYZ.get(functionCentre.get(i) - 1), (Double)objects.get(i), 2, 3, 3));  // create an fyzz-type function
                    } else if (functionTypes.get(i) == 20) {
                        atomicFunctions.add(new fAF(XYZ.get(functionCentre.get(i) - 1), (Double)objects.get(i), 3, 3, 3));  // create an fzzz-type function
                    }
                }
            }
            line = br.readLine();  // read next line
        }
        //System.out.println(numPrimitives + " primitives and " + functionTypes.size() + " function exponents");
        // read in molecular orbitals
        //DO 120 I = 1,NMO
        //READ (IWFN,106) CO(IP+I),CO(IEORB+I)
        //READ (IWFN,107) (CO(K+J),J=1,NPRIMS)
        for (int i = 0; i < numOrbitals; ++i) {
            objects = FortranFormat.read(line,"(35X,F12.8,15X,F12.8)");
            orbitalOccupation.add((Double)objects.get(0));  // store away orbital occupation
            orbitalEnergy.add((Double)objects.get(1));  // store away orbital energy
            List<Double> orbitalCoeff = new ArrayList<Double>();  // temp ArrayList of orbital exponents
            line = br.readLine();  // read next line
            while (!(line.startsWith("MO")) && !(line.startsWith("END"))) {
                objects = FortranFormat.read(line,"(5E16.8)");
                for (int j = 0; j < 5; ++j) {  // 5 objects per line
                    if (objects.get(j) != null) {
                        orbitalCoeff.add((Double)objects.get(j));
                    }
                }
                line = br.readLine();  // read next line
            }
	    for (int j = 0; j < numPrimitives; ++j) {  
                atomicFunctions.get(j).coeff.add(orbitalCoeff.get(j)); // finally, store away this MO's orbital coefficients into the AF objects
	    }
        }
        line = br.readLine();  // read next line
        objects = FortranFormat.read(line,"(17X,F20.12,18X,F13.8)");
        totE = (Double)objects.get(0);
        gamma = (Double)objects.get(1);  // virial ratio
    }

    public String toString() {
        return titleWFN + " TOTAL ENERGY = " + totE ;
    }

    // determine the charge density, rho, at Centre p
    public Double getRho(Centre p) {
	 	double[] psi = new double[numOrbitals];
		Double rho = 0.0;
      for (AF af : atomicFunctions) {  // loop over primitives
       	double bf = af.getMOValue(p);  // get exponent at p
	    		for (int i = 0; i < numOrbitals; ++i) { // loop over molecular orbitals
            	psi[i] += bf*af.coeff.get(i);  // determine psi(MO)
	    		}
			}
	 	   // now sum up orbital contributions
			for (int i = 0; i < numOrbitals; ++i) { // loop over molecular orbitals
         	rho += psi[i]*psi[i]*orbitalOccupation.get(i);
			}
	
        return rho;
	}

	public double[] getGradRho(Centre p) {
		double[] delRHo = new double[3];  // gradient vector
		double[] gx = new double[numOrbitals];
		double[] gy = new double[numOrbitals];
		double[] gz = new double[numOrbitals];
		for (AF af : atomicFunctions) {  // loop over primitives
       	double bf = af.getMOValue(p);  // get exponent at p
	    		for (int i = 0; i < numOrbitals; ++i) { // loop over molecular orbitals
            	psi[i] += bf*af.coeff.get(i);  // determine psi(MO)
	    		}
			}
	 	   // now sum up orbital contributions
			//      DO 40 I = 1,NMO
 		   //     P0 = CO(IP+I)
   		//     P1 = P0*CO(IPSI+I)
  			//      W(1) = W(1) + P1*CO(IGX+I)
  			//      W(2) = W(2) + P1*CO(IGY+I)
			//40      W(3) = W(3) + P1*CO(IGZ+I)	
			for (int i = 0; i < numOrbitals; ++i) { // loop over molecular orbitals
         	delRho[0] += gx[i]*psi[i]*orbitalOccupation.get(i);
         	delRho[1] += gy[i]*psi[i]*orbitalOccupation.get(i);
         	delRho[2] += gz[i]*psi[i]*orbitalOccupation.get(i);
			}
	
        return rho;

	}
}
